﻿namespace SpikingAnalyticsLib

open System
open System.IO
open SpikingAnalyticsLib.GeneralUtilities

//=================================================================================================

/// A response frame: wraps a map of firing time offsets for each neuron collected over a fixed interval in response to a specified stimulus
type Frame(offsetMap:Map<int, int list>) =

//---------------------------------------------------------------------------------------------------------
// Properties

    // Get the firing time offsets for each neuron
    member this.OffsetMap = offsetMap

//---------------------------------------------------------------------------------------------------------
// Methods

    /// Save the frame to the specified stream
    member this.Save(sw:StreamWriter) =
        offsetMap
        |> Map.toSeq
        |> Seq.iter (fun (nindex, firingTimes) ->
            let line = String.Join(" ", (nindex :: firingTimes))
            sw.WriteLine(line)
        )

    /// Save the frame to the specified file
    member this.Save(filePath:string) =
        use sw = File.CreateText(filePath)
        this.Save(sw)

    /// Load a frame from the specified stream
    static member Load(sr:StreamReader) =
        let offsetMap =
            seq {
                // read a sequence of strings until an empty line (or EOF) is encountered
                while (sr.Peek() >= 0) do
                    yield sr.ReadLine()
            }
            |> Seq.takeWhile (fun line -> not (String.IsNullOrEmpty line))
            |> Seq.map (fun line ->
                let data = IOUtilities.ParseStringToList<int>(line)
                data.Head, data.Tail
            )
            |> Map.ofSeq
        new Frame(offsetMap)

    /// Load a frame from the specified file
    static member Load(filePath:string) =
        use sr = File.OpenText(filePath)
        Frame.Load(sr)

//=================================================================================================

/// A set of response frames
type FrameSet (frames:Frame []) =

    // Process each frame
    static let ProcessFrames selectedNeurons totalNeurons frameSize runSeconds firingEvents =

        let selectionMask = CreateFilter selectedNeurons totalNeurons
        let frameArray:Frame option [] = Array.init (runSeconds * 1000 / frameSize) (fun index -> None)

        let finalFrameIndex, finalFrameMap =
            firingEvents
            |> Seq.filter (fun (_, nindex, _) -> selectionMask.[nindex])                        // process only selected neurons
            |> Seq.map (fun (time, nindex, _) -> time / frameSize, time % frameSize, nindex)    // remap time into each frame of size 'frameSize'
            |> Seq.fold (fun (state:int * Map<int,int list>) (frameIndex, time, nindex) ->
                    let currentFrame, map = state
                    if frameIndex = currentFrame then
                        // produce a list of reponse times for each neuron in the current frame
                        let times =
                            if map.ContainsKey(nindex) then
                                map.[nindex]
                            else
                                []
                        currentFrame, Map.add nindex (time :: times) map
                    else
                        // process the previous frame
                        let frame = new Frame(map)
                        frameArray.[currentFrame] <- Some(frame)

                        // start a new frame
                        frameIndex, Map.add nindex ([ time ]) Map.empty
                ) (0, Map.empty)

        frameArray.[finalFrameIndex] <- Some(new Frame(finalFrameMap))

        //assert Array.forall (fun (frame:Frame option) -> frame.IsSome) frameArray

        frameArray
        |> Array.mapi (fun index frame ->
                if frame.IsSome then
                    frame.Value
                else
                    // add an empty frame
                    printfn "Warning: empty frame at index %d" index
                    new Frame(Map.empty)
            )

//---------------------------------------------------------------------------------------------------------
// Constructors

    new(selectedNeurons, totalNeurons, frameSize, runSeconds, firingData:seq<int * int * SpikingNeuronLib.EventLabel>) =
        let frameArray = ProcessFrames selectedNeurons totalNeurons frameSize runSeconds firingData
        new FrameSet(frameArray)

//---------------------------------------------------------------------------------------------------------
// Properties

    /// Get the number of Frame instances
    member this.Length = frames.Length

    /// Retrieve a specific Frame instance
    member this.Item
        with get(index) = frames.[index]
        //and set index value = frames.[index] <- value

//---------------------------------------------------------------------------------------------------------
// Methods

    /// Return a sequence of Frame instances
    member this.ToEnumerable() = Seq.init frames.Length (fun index -> frames.[index])

    /// Save the frame set to the specified file
    member this.Save(filePath:string) =
        use sw = File.CreateText(filePath)
        for frame in frames do
            frame.Save(sw)
            sw.WriteLine()      // write an empty line at the end of each frame

    /// Load a frame set from the specified file
    static member Load(filePath:string) =
        use sr = File.OpenText(filePath)
        let frameArray =
            seq {
                while (sr.Peek() >= 0) do
                    yield Frame.Load(sr)
            }
            |> Seq.toArray
        new FrameSet(frameArray)
